extern crate nalgebra as na;
use na::*;
use nalgebra::{
    base::allocator::Allocator, DefaultAllocator, Dim, DimName, DimNameAdd, MatrixMN, RealField,
    Scalar,
};
use rand::prelude::*;
use std::ops::Add;

extern crate docify;
use nalgebra;
use nalgebra::DimMinimum;

/// https://github.com/rsasaki0109/rust_robotics/blob/main/src/bin/ekf.rs
/// 更新 x_hat/status,P/process matrix    
///  x_hat: Matrix2x1<f64> Initial state estimate
///   F: Matrix2<f64>,       // State transition matrix ,扩展卡尔曼JF
///  H: Matrix1x2<f64>,     // Observation matrix,扩展卡尔曼JH
///  B: Matrix2<f64>,       //  control var matrix , u is contral var
///  K: Matrix2x1<f64>,     //卡尔曼增益
/// Q: Matrix2<f64>,       // Process noise covariance matrix
///  R: Matrix1<f64>,       // Measurement noise covariance matrix
///  P: Matrix2<f64>,       // Initial error covariance
pub struct Kalman {
    pub x_hat: Matrix2x1<f64>, // Initial state estimate
    pub F: Matrix2<f64>,       // State transition matrix ,扩展卡尔曼JF
    pub H: Matrix1x2<f64>,     // Observation matrix,扩展卡尔曼JH
    pub B: Matrix2<f64>,       //  control var matrix , u is contral var
    pub K: Matrix2x1<f64>,     //卡尔曼增益
    pub Q: Matrix2<f64>,       // Process noise covariance matrix
    pub R: Matrix1<f64>,       // Measurement noise covariance matrix
    pub P: Matrix2<f64>,       // Initial error covariance
}

#[doc = docify::embed!("src/robotics/kalman.rs", test_kalman)]
impl Kalman {
    pub fn predict(&mut self, u: Matrix2x1<f64>) {
        //  扩展kalman对F求雅可比jF      jF = jacob_f(xEst, u)
        self.x_hat = self.F * self.x_hat + self.B * u;
        self.P = self.F * self.P * self.F.transpose() + self.Q;
    }
    // 更新 x_hat,K,P
    pub fn update(&mut self, z: f64) {
        //  扩展kalman对H求雅可比jH
        let _inv = self.H * self.P * self.H.transpose() + self.R;
        let inv = _inv.try_inverse().unwrap();
        //  扩展kalman对H求雅可比jH
        self.K = self.P * self.H.transpose() * inv;
        let y = self.H * self.x_hat; //observation model
        let dif = z - y[(0, 0)];
        self.x_hat = self.x_hat + self.K * dif;
        let eye = Matrix2::<f64>::identity();
        self.P = (eye - self.K * self.H) * self.P;
    }
}

#[test]
#[docify::export]
fn test_kalman() {
    let F = Matrix2::new(1., 0., 1., 1.);
    let B = Matrix2::new(1., 0., 1., 1.);
    let H = Matrix1x2::new(1., 0.);
    let Q = Matrix2::<f64>::identity() * 0.01;
    let R = Matrix1::new(0.1);
    let x_hat = Matrix2x1::<f64>::zeros();
    let P = Matrix2::<f64>::identity() * 0.1;
    let u = Matrix2x1::<f64>::zeros();
    let k = Matrix2x1::<f64>::zeros();

    let mut kal = Kalman {
        x_hat: x_hat,
        F: F,
        H: H,
        B: B,
        K: k,
        Q: Q,
        R: R,
        P: P,
    };

    let mut rng = rand::thread_rng();
    //  true state you never know
    let true_state: Vec<f64> = (0..100).map(|_| rng.gen::<f64>() / 10.0).collect();
    let measurements: Vec<f64> = true_state
        .iter()
        .map(|v| v + rng.gen::<f64>() / 10.0)
        .collect();
    for z in measurements {
        // should + B*u ,but not        predict func
        &kal.predict(u);
        &kal.update(z);
    }
}